#!/usr/bin/env python3
"""
测试Coordinator服务器的功能
验证ADD、DEL、SET事件处理和liveness接口
"""

import asyncio
import json
import time
import sys
import os
from datetime import datetime, timezone

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

import httpx
from motor.resources.http_msg_spec import InsEventMsg, EventType
from motor.resources.instance import Instance, ParallelConfig
from motor.resources.endpoint import Endpoint, DeviceInfo, EndpointStatus

# 测试配置
SERVER_URL = "http://localhost:8080"
TEST_TIMEOUT = 30

async def test_coordinator_server():
    """测试Coordinator服务器功能"""
    print("🚀 开始测试Coordinator服务器...")
    
    async with httpx.AsyncClient(timeout=TEST_TIMEOUT) as client:
        # 1. 测试根路径
        print("\n1. 测试根路径...")
        try:
            response = await client.get(f"{SERVER_URL}/")
            print(f"✅ 根路径响应: {response.status_code}")
            print(f"   内容: {response.json()}")
        except Exception as e:
            print(f"❌ 根路径测试失败: {e}")
            return False
        
        # 2. 测试liveness接口
        print("\n2. 测试liveness接口...")
        try:
            response = await client.get(f"{SERVER_URL}/liveness")
            print(f"✅ Liveness响应: {response.status_code}")
            print(f"   内容: {response.json()}")
        except Exception as e:
            print(f"❌ Liveness测试失败: {e}")
            return False
        
        # 3. 测试heartbeat接口
        print("\n3. 测试heartbeat接口...")
        try:
            response = await client.post(f"{SERVER_URL}/coordinator/heartbeat", json={"status": "normal"})
            print(f"✅ Heartbeat响应: {response.status_code}")
            print(f"   内容: {response.json()}")
        except Exception as e:
            print(f"❌ Heartbeat测试失败: {e}")
            return False
        
        # 4. 创建测试实例
        print("\n4. 创建测试实例...")
        test_instance = Instance(
            job_name="test-job-1",
            id=1,
            role="worker",
            endpoints={
                "127.0.0.1": {
                    1: Endpoint(
                        id=1,
                        ip="127.0.0.1",
                        port="8080",
                        status=EndpointStatus.NORMAL,
                        device_infos=[
                            DeviceInfo(
                                device_id="0",
                                device_ip="127.0.0.1",
                                rank_id="0"
                            )
                        ]
                    )
                }
            },
            parallel_config=ParallelConfig(
                dp=1,
                cp=1,
                tp=1
            )
        )
        
        # 5. 测试ADD事件
        print("\n5. 测试ADD事件...")
        try:
            add_event = InsEventMsg(
                event=EventType.ADD,
                instances=[test_instance]
            )
            response = await client.post(f"{SERVER_URL}/coordinator/active", json=add_event.model_dump())
            print(f"✅ ADD事件响应: {response.status_code}")
            print(f"   内容: {response.json()}")
        except Exception as e:
            print(f"❌ ADD事件测试失败: {e}")
            return False
        
        # 6. 检查实例列表
        print("\n6. 检查实例列表...")
        try:
            response = await client.get(f"{SERVER_URL}/instances")
            print(f"✅ 实例列表响应: {response.status_code}")
            instances_data = response.json()
            print(f"   实例数量: {instances_data['count']}")
            print(f"   实例列表: {list(instances_data['instances'].keys())}")
        except Exception as e:
            print(f"❌ 实例列表测试失败: {e}")
            return False
        
        # 7. 测试SET事件（全量更新）
        print("\n7. 测试SET事件...")
        try:
            # 创建多个实例
            instances = [
                Instance(
                    job_name=f"test-job-{i}",
                    id=i,
                    role="worker",
                    endpoints={
                        "127.0.0.1": {
                            i: Endpoint(
                                id=i,
                                ip="127.0.0.1",
                                port=f"808{i}",
                                status=EndpointStatus.NORMAL,
                                device_infos=[
                                    DeviceInfo(
                                        device_id=str(i-1),
                                        device_ip="127.0.0.1",
                                        rank_id=str(i-1)
                                    )
                                ]
                            )
                        }
                    },
                    parallel_config=ParallelConfig(
                        dp=1,
                        cp=1,
                        tp=1
                    )
                ) for i in range(1, 4)
            ]
            
            set_event = InsEventMsg(
                event=EventType.SET,
                instances=instances
            )
            response = await client.post(f"{SERVER_URL}/coordinator/active", json=set_event.model_dump())
            print(f"✅ SET事件响应: {response.status_code}")
            print(f"   内容: {response.json()}")
        except Exception as e:
            print(f"❌ SET事件测试失败: {e}")
            return False
        
        # 8. 再次检查实例列表
        print("\n8. 再次检查实例列表...")
        try:
            response = await client.get(f"{SERVER_URL}/instances")
            instances_data = response.json()
            print(f"✅ 更新后实例数量: {instances_data['count']}")
            print(f"   实例列表: {list(instances_data['instances'].keys())}")
        except Exception as e:
            print(f"❌ 实例列表检查失败: {e}")
            return False
        
        # 9. 测试DEL事件
        print("\n9. 测试DEL事件...")
        try:
            del_event = InsEventMsg(
                event=EventType.DEL,
                instances=[instances[0]]  # 删除第一个实例
            )
            response = await client.post(f"{SERVER_URL}/coordinator/active", json=del_event.model_dump())
            print(f"✅ DEL事件响应: {response.status_code}")
            print(f"   内容: {response.json()}")
        except Exception as e:
            print(f"❌ DEL事件测试失败: {e}")
            return False
        
        # 10. 最终检查实例列表
        print("\n10. 最终检查实例列表...")
        try:
            response = await client.get(f"{SERVER_URL}/instances")
            instances_data = response.json()
            print(f"✅ 删除后实例数量: {instances_data['count']}")
            print(f"   实例列表: {list(instances_data['instances'].keys())}")
        except Exception as e:
            print(f"❌ 最终实例列表检查失败: {e}")
            return False
        
        # 11. 测试完整状态
        print("\n11. 测试完整状态...")
        try:
            response = await client.get(f"{SERVER_URL}/status")
            status_data = response.json()
            print(f"✅ 状态响应: {response.status_code}")
            print(f"   协调器状态: {status_data['coordinator_status']}")
            print(f"   事件历史数量: {len(status_data['event_history'])}")
        except Exception as e:
            print(f"❌ 状态测试失败: {e}")
            return False
        
        # 12. 最终liveness检查
        print("\n12. 最终liveness检查...")
        try:
            response = await client.get(f"{SERVER_URL}/liveness")
            liveness_data = response.json()
            print(f"✅ 最终liveness: {response.status_code}")
            print(f"   服务状态: {liveness_data['status']}")
            print(f"   实例数量: {liveness_data['instance_count']}")
        except Exception as e:
            print(f"❌ 最终liveness检查失败: {e}")
            return False
    
    print("\n🎉 所有测试完成！")
    return True

async def main():
    """主函数"""
    print("=" * 60)
    print("Motor Coordinator Server 测试脚本")
    print("=" * 60)
    
    # 等待服务器启动
    print("⏳ 等待服务器启动...")
    await asyncio.sleep(2)
    
    # 运行测试
    success = await test_coordinator_server()
    
    if success:
        print("\n✅ 所有测试通过！")
    else:
        print("\n❌ 部分测试失败！")
    
    print("=" * 60)

if __name__ == "__main__":
    asyncio.run(main())
